#!/bin/bash
# 版本: 1.0
# 作者: huangsheng
# 兼容系统版本: ubuntu
# 验证: 未验证

# 定义变量
software_dir="stable-diffusion"
venv_name="amd${software_dir}"
ONEGPU_RUN_SD="onegpu_run_sd.sh"
MGPU_INSTALL="mgpu_sd_install.sh"
SCRIPT_PATH=$(realpath "$0")

# 定义显卡型号与HSA_OVERRIDE_GFX_VERSION的对应关系
declare -A version_map=(
    ["1"]="RX 7900 XTX, RX 7900 XT, RX 7900 GRE, PRO W7900, PRO W7800"
    ["2"]="RX 7800 XT, RX 7700 XT, PRO W7700"
    ["3"]="RX 7600"
    ["4"]="RX 6900 XT, RX 6950 XT, RX 6800 XT, RX 6800, PRO V620, PRO W6800"
    ["5"]="RX 6700 XT, RX 6750 XT, RX 6700"
    ["6"]="RX 6600 XT, RX 6600, RX 6650 XT, PRO W6600"
    ["7"]="Pro W5500"
    ["8"]="VII, Pro VII"
)

# HSA_OVERRIDE_GFX_VERSION对应关系
declare -A gfx_version_map=(
    ["1"]="11.0.0"
    ["2"]="11.0.1"
    ["3"]="11.0.2"
    ["4"]="10.3.0"
    ["5"]="10.3.1"
    ["6"]="10.3.2"
    ["7"]="10.1.2"
    ["8"]="9.0.6"
)

# 显示可用的HSA_OVERRIDE_GFX_VERSION选项
echo "请选择对应的显卡版本编号："
for i in {1..8}; do
    echo "$i: ${version_map[$i]}"
done

# 读取用户输入
read -p "输入显卡版本编号，如果显卡不在上面的列表中可以输入相近的型号的数字（或输入 'N' 退出）: " version_choice

# 验证用户输入并查找对应的版本号
if [[ "$version_choice" == "N" || "$version_choice" == "n" ]]; then
    echo "退出脚本。"
    exit 0
elif [[ -v version_map[$version_choice] ]]; then
    found_version=${gfx_version_map[$version_choice]}
    echo "选择的HSA_OVERRIDE_GFX_VERSION: $found_version"
else
    echo "无效的显卡版本编号。"
    echo "请重新运行脚本，并输入有效的编号。"
    exit 1
fi

# 安装必要的软件包
sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0 -y

# 下载文件
[ ! -d "$software_dir" ] && git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git $software_dir
cd $software_dir
[ ! -d "$venv_name" ] && python3 -m venv $venv_name
source $venv_name/bin/activate

# 安装依赖
pip3 install --upgrade pip
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0
pip3 install -r requirements.txt

# 建立启动脚本
if [ ! -f "$ONEGPU_RUN_SD" ]; then
    cat > $ONEGPU_RUN_SD <<EOF
#!/bin/bash
source $venv_name/bin/activate
export HSA_OVERRIDE_GFX_VERSION=$found_version
export HIP_VISIBLE_DEVICES=0
python3 launch.py --lowvram --precision full --no-half --autolaunch --listen --theme dark --skip-torch-cuda-test
EOF
    chmod +x $ONEGPU_RUN_SD
fi

# 提示用户是否安装多显卡服务
read -p "是否安装多显卡服务？(Y/N): " install_multi_gpu

if [[ "$install_multi_gpu" == "Y" || "$install_multi_gpu" == "y" ]]; then
    # 下载多显卡服务脚本
    wget -O $MGPU_INSTALL https://gitee.com/huangsheng32/AMDGPU_linux/raw/master/stable-diffusion/mgpu_sd_install.sh
    chmod +x $MGPU_INSTALL

    # 执行多显卡服务脚本
    bash $MGPU_INSTALL
else
    echo "未安装多显卡服务。"
fi

# 删除脚本文件自身
sudo rm -- "$SCRIPT_PATH"

echo "安装完成"
